import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from libs.loss import iou_loss, bineary_iou_loss
from libs.utils import save_sample_imgs, showimg, showmask, AverageMeter,batch_end_log, epoch_end_log, make_ckpt_name
from tqdm import tqdm
import importlib
from trainer.BaseTrainer import BaseTrainer
import torch.nn as nn
import os

class Trainer(BaseTrainer):
    def __init__(self, config, train_data_loader, test_data_loader):
        super(Trainer, self).__init__(config, train_data_loader, test_data_loader)


    def build_model(self):
        module = importlib.import_module("nets.hairnet")
        self.net = module.get_seg_model()
        self.net.to(self.device)


    def eval(self):
        print("star to test...")
        self.net.eval()
        with torch.no_grad():
            eval_loss_monitor = AverageMeter()
            eval_aiou_monitor = AverageMeter()
            eval_biou_monitor = AverageMeter()

            eval_loss_monitor.reset()
            eval_aiou_monitor.reset()
            eval_biou_monitor.reset()

            for samples in tqdm(self.test_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask_gt'].to(self.device)
                alpha_gt = samples['blurmask_gt'].to(self.device)

                mask_pred = self.net(image)

                eval_loss_monitor.update(self.lossfunc(mask_pred, mask_gt).item())
                eval_biou_monitor.update(bineary_iou_loss(mask_pred, mask_gt).item())
                eval_aiou_monitor.update(iou_loss(mask_pred, alpha_gt).item())


            avg_loss = eval_loss_monitor.avg
            avg_aiou = eval_aiou_monitor.avg
            avg_biou = eval_biou_monitor.avg

            print("test phase | loss: %.4f amask_iou: %.4f amask_iou: %.4f" %(avg_loss, avg_aiou, avg_biou))

        return avg_loss, avg_aiou, avg_biou



    def train(self):

        optimizer = torch.optim.Adam(self.net.parameters(), lr=self.config.lr, eps=1e-7)

        lr_scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=4, verbose=True, threshold=0.015)

        train_loss_monitor = AverageMeter()
        biou_loss_monitor = AverageMeter()
        aiou_loss_monitor = AverageMeter()

        self.lossfunc = nn.BCELoss()

        for epoch in range(self.config.max_epoch):
            self.net.train()
            train_loss_monitor.reset()
            aiou_loss_monitor.reset()
            biou_loss_monitor.reset()

            for step, samples in enumerate(self.train_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask_gt'].to(self.device)
                alpha_gt = samples["blurmask_gt"].to(self.device)


                mask_pred = self.net(image)
                self.net.zero_grad()

                train_loss = self.lossfunc(mask_pred, mask_gt)
                train_loss_value = train_loss.item()
                train_loss_monitor.update(train_loss_value)


                biou_value = bineary_iou_loss(mask_pred, mask_gt).item()
                biou_loss_monitor.update(biou_value)

                aiou_value = iou_loss(mask_pred, alpha_gt).item()
                aiou_loss_monitor.update(aiou_value)

                train_loss.backward()
                optimizer.step()

                batch_end_log(log_intervel=self.config.log_intervel, max_epoch=self.config.max_epoch,
                              epoch=epoch,
                              step=step, max_step=self.train_image_len, lr=optimizer.param_groups[0]['lr'],
                              train_loss=train_loss_value, aiou=aiou_value, biou=biou_value)


                if step % self.config.sample_intervel == 0:
                    save_sample_imgs(samples, {"mask_pred": mask_pred, "dyemask": mask_pred}, self.config.sample_dir, epoch, step)


            epoch_end_log(phase="train", train_loss=train_loss_monitor.avg, aiou=aiou_loss_monitor.avg, biou=biou_loss_monitor.avg)

            lr_scheduler.step(train_loss_monitor.avg)


            test_loss, test_aiou, test_biou = self.eval()
            ckpt_name = make_ckpt_name(trainer=self.config.trainer, train_loss=train_loss_monitor.avg, aiou=aiou_loss_monitor.avg,
                                        suffix="pth")
            torch.save(self.net.state_dict(), os.path.join(self.config.checkpoint_dir, ckpt_name))